"""
数据血缘分析服务
"""

import networkx as nx
from typing import Dict, List, Optional, Tuple, Any
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
import uuid
import time


from app.models.lineage import DataLineage, LineageNode, LineageEdge
from app.schemas.lineage import DataLineageSchema, LineageNodeSchema, LineageEdgeSchema
from app.services.sql_parser import SQLParserService
from app.core.logger import logger


class LineageAnalyzerService:
    """数据血缘分析服务"""
    
    def __init__(self):
        self.sql_parser = SQLParserService()
        self.graph = nx.DiGraph()
    
    async def analyze_task_lineage(
        self, 
        db: AsyncSession, 
        task_id: int, 
        force_refresh: bool = False
    ) -> Optional[DataLineageSchema]:
        """
        分析单个任务的数据血缘
        
        Args:
            db: 数据库会话
            task_id: 任务ID
            force_refresh: 是否强制刷新
            
        Returns:
            数据血缘信息
        """
        start_time = time.time()
        
        try:
            # 获取任务信息
            task = await self._get_task_by_id(db, task_id)
            if not task:
                logger.error(f"任务不存在: {task_id}")
                return None
            
            if not task.statement:
                logger.warning(f"任务 {task_id} 没有SQL语句")
                return None
            
            # 检查是否已存在血缘信息
            if not force_refresh:
                existing_lineage = await self._get_existing_lineage(db, task_id)
                if existing_lineage:
                    logger.info(f"使用缓存的血缘信息: {task_id}")
                    return existing_lineage
            
            # 解析SQL语句
            parsed_result = self.sql_parser.parse_flink_sql(task.statement)
            
            # 构建血缘图
            lineage_data = self._build_lineage_graph(
                task_id=task_id,
                task_name=task.name,
                parsed_result=parsed_result
            )
            
            # 保存到数据库
            await self._save_lineage_to_db(db, lineage_data)
            
            processing_time = time.time() - start_time
            logger.info(f"血缘分析完成: {task_id}, 耗时: {processing_time:.2f}s")
            
            return lineage_data
            
        except Exception as e:
            logger.error(f"血缘分析失败: {task_id}, 错误: {str(e)}")
            return None
    
    async def _get_task_by_id(self, db: AsyncSession, task_id: int) -> Optional[dict]:
        """根据ID获取任务（从解析任务表）"""
        from sqlalchemy import text
        result = await db.execute(text("""
            SELECT id, task_name as name, sql_content as statement
            FROM dp_parse_tasks
            WHERE id = :task_id AND parse_status = 'SUCCESS'
        """), {"task_id": task_id})
        row = result.fetchone()
        if row:
            return {"id": row.id, "name": row.name, "statement": row.statement}
        return None
    
    async def _get_existing_lineage(self, db: AsyncSession, task_id: int) -> Optional[DataLineageSchema]:
        """获取已存在的血缘信息"""
        result = await db.execute(
            select(DataLineage).where(DataLineage.task_id == task_id)
        )
        lineage = result.scalar_one_or_none()
        
        if lineage:
            # 获取节点和边
            nodes_result = await db.execute(
                select(LineageNode).where(LineageNode.lineage_id == lineage.id)
            )
            edges_result = await db.execute(
                select(LineageEdge).where(LineageEdge.lineage_id == lineage.id)
            )
            
            nodes = [self._convert_node_to_schema(node) for node in nodes_result.scalars()]
            edges = [self._convert_edge_to_schema(edge) for edge in edges_result.scalars()]
            
            return DataLineageSchema(
                id=lineage.id,
                task_id=lineage.task_id,
                task_name=lineage.task_name,
                nodes=nodes,
                edges=edges,
                node_count=len(nodes),
                edge_count=len(edges),
                status=lineage.status,
                created_at=lineage.created_at,
                updated_at=lineage.updated_at
            )
        
        return None
    
    def _build_lineage_graph(
        self, 
        task_id: int, 
        task_name: str, 
        parsed_result: Dict[str, List[str]]
    ) -> DataLineageSchema:
        """构建血缘图"""
        
        # 重置图
        self.graph.clear()
        
        nodes = []
        edges = []
        
        # 创建节点
        all_tables = set()
        all_tables.update(parsed_result['source_tables'])
        all_tables.update(parsed_result['target_tables'])
        all_tables.update(parsed_result['temp_tables'])
        all_tables.update(parsed_result['views'])
        
        for table in all_tables:
            node_type = self._determine_node_type(table, parsed_result)
            database, table_name = self._split_table_name(table)
            
            node = LineageNodeSchema(
                id=f"node_{uuid.uuid4().hex[:8]}",
                type=node_type,
                database=database,
                table=table_name,
                full_name=table,
                columns=[],
                properties={"task_id": task_id},
                position={"x": 0, "y": 0}
            )
            nodes.append(node)
            self.graph.add_node(node.id, **node.dict())
        
        # 创建边
        dependencies = self.sql_parser.extract_table_dependencies("")
        
        # 从解析结果创建边
        node_map = {node.full_name: node.id for node in nodes}
        
        # 只有在INSERT INTO ... SELECT FROM 语句中才建立源表到目标表的边
        # 避免创建双向或无意义的连接
        if parsed_result['source_tables'] and parsed_result['target_tables']:
            # 对于每个目标表，只连接到真正的源表
            for target_table in parsed_result['target_tables']:
                for source_table in parsed_result['source_tables']:
                    if (source_table in node_map and target_table in node_map and
                        source_table != target_table):  # 避免自循环
                        edge = LineageEdgeSchema(
                            id=f"edge_{uuid.uuid4().hex[:8]}",
                            source=node_map[source_table],
                            target=node_map[target_table],
                            type="data_transformation",
                            properties={"task_id": task_id, "sql_type": "INSERT_SELECT"}
                        )
                        edges.append(edge)
                        self.graph.add_edge(edge.source, edge.target, **edge.dict())
        
        # 源表到视图的边
        for view in parsed_result['views']:
            for source_table in parsed_result['source_tables']:
                if source_table in node_map and view in node_map:
                    edge = LineageEdgeSchema(
                        id=f"edge_{uuid.uuid4().hex[:8]}",
                        source=node_map[source_table],
                        target=node_map[view],
                        type="view_dependency",
                        properties={"task_id": task_id}
                    )
                    edges.append(edge)
                    self.graph.add_edge(edge.source, edge.target, **edge.dict())
        
        # 计算布局
        self._calculate_layout(nodes)
        
        return DataLineageSchema(
            task_id=task_id,
            task_name=task_name,
            nodes=nodes,
            edges=edges,
            node_count=len(nodes),
            edge_count=len(edges),
            status="active"
        )
    
    def _determine_node_type(self, table: str, parsed_result: Dict[str, List[str]]) -> str:
        """确定节点类型"""
        if table in parsed_result['views']:
            return "view"
        elif table in parsed_result['temp_tables']:
            return "temp"
        else:
            return "table"
    
    def _split_table_name(self, full_name: str) -> Tuple[Optional[str], str]:
        """分割表名"""
        if '.' in full_name:
            parts = full_name.split('.')
            return parts[0], parts[1]
        return None, full_name
    
    def _calculate_layout(self, nodes: List[LineageNodeSchema]):
        """计算节点布局"""
        if not nodes:
            return
        
        try:
            # 使用层次布局
            if self.graph.number_of_nodes() > 0:
                pos = nx.spring_layout(self.graph, k=3, iterations=50)
                
                for node in nodes:
                    if node.id in pos:
                        node.position = {
                            "x": int(pos[node.id][0] * 300),
                            "y": int(pos[node.id][1] * 300)
                        }
        except Exception as e:
            logger.warning(f"布局计算失败: {str(e)}")
            # 使用简单的网格布局
            for i, node in enumerate(nodes):
                node.position = {
                    "x": (i % 5) * 200,
                    "y": (i // 5) * 150
                }
    
    async def _save_lineage_to_db(self, db: AsyncSession, lineage_data: DataLineageSchema):
        """保存血缘信息到数据库"""
        try:
            # 删除已存在的血缘信息
            await db.execute(
                select(DataLineage).where(DataLineage.task_id == lineage_data.task_id)
            )
            existing = await db.execute(
                select(DataLineage).where(DataLineage.task_id == lineage_data.task_id)
            )
            existing_lineage = existing.scalar_one_or_none()
            
            if existing_lineage:
                await db.delete(existing_lineage)
            
            # 创建新的血缘记录
            lineage = DataLineage(
                task_id=lineage_data.task_id,
                task_name=lineage_data.task_name,
                lineage_data=lineage_data.dict(),
                node_count=lineage_data.node_count,
                edge_count=lineage_data.edge_count,
                status=lineage_data.status
            )
            
            db.add(lineage)
            await db.flush()
            
            # 保存节点
            for node_data in lineage_data.nodes:
                node = LineageNode(
                    lineage_id=lineage.id,
                    node_id=node_data.id,
                    node_type=node_data.type,
                    database_name=node_data.database,
                    table_name=node_data.table,
                    full_name=node_data.full_name,
                    columns=node_data.columns,
                    properties=node_data.properties,
                    position_x=node_data.position.get("x", 0) if node_data.position else 0,
                    position_y=node_data.position.get("y", 0) if node_data.position else 0
                )
                db.add(node)
            
            # 保存边
            for edge_data in lineage_data.edges:
                edge = LineageEdge(
                    lineage_id=lineage.id,
                    edge_id=edge_data.id,
                    source_node_id=edge_data.source,
                    target_node_id=edge_data.target,
                    edge_type=edge_data.type,
                    column_mapping=edge_data.column_mapping,
                    properties=edge_data.properties
                )
                db.add(edge)
            
            await db.commit()
            logger.info(f"血缘信息已保存: task_id={lineage_data.task_id}")
            
        except Exception as e:
            await db.rollback()
            logger.error(f"保存血缘信息失败: {str(e)}")
            raise
    
    def _convert_node_to_schema(self, node: LineageNode) -> LineageNodeSchema:
        """转换节点模型到Schema"""
        return LineageNodeSchema(
            id=node.node_id,
            type=node.node_type,
            database=node.database_name,
            table=node.table_name,
            full_name=node.full_name,
            columns=node.columns or [],
            properties=node.properties or {},
            position={"x": node.position_x, "y": node.position_y}
        )
    
    def _convert_edge_to_schema(self, edge: LineageEdge) -> LineageEdgeSchema:
        """转换边模型到Schema"""
        return LineageEdgeSchema(
            id=edge.edge_id,
            source=edge.source_node_id,
            target=edge.target_node_id,
            type=edge.edge_type,
            column_mapping=edge.column_mapping or {},
            properties=edge.properties or {}
        )
